#pragma once

#include "file_comm.h"
#include "xxhash.h"
#include "Socket.h"
#include <ios>
#include <fstream>
#include <mutex>
#include <filesystem>


class FileReader {
public:
    explicit FileReader([[maybe_unused]] const char *path) :
            file(std::filesystem::path(path), std::ios::binary) {
        file.seekg(0, std::ios::end);
        file_size = file.tellg();
        file.seekg(0, std::ios::beg);
    }

    int getCount() const {
        return int(file_size / CHUNK_SIZE) + bool(file_size % CHUNK_SIZE);
    }

    Chunk getNext() {
        std::unique_lock<std::mutex> lock(mtx);
        if (current_pos >= file_size) {
            return {};
        }
        Chunk chunk{
                .head = {},
                .data = std::make_unique<char[]>(CHUNK_SIZE),
        };
        chunk.head.no = current_no++;
        file.read(chunk.data.get(), CHUNK_SIZE);
        chunk.head.size = (int) file.gcount();
        chunk.head.hash = XXHash::xxh32(chunk.data.get(), chunk.head.size, 0);
        current_pos += chunk.head.size;
        return chunk;
    }


private:
    std::ifstream file;
    size_t file_size = 0;
    size_t current_pos = 0;
    int current_no = 1;
    std::mutex mtx;
};

class FileSender {
public:
    FileSender(const char *path, const char *host, int port) :
            reader(path), host(host), port(port) {}

    void sendFile() {
        sendCollect();

        std::thread ts[6];
        for (auto &t: ts) {
            t = std::thread(&FileSender::sendChunkWorker, this);
        }
        for (auto &t: ts) {
            t.join();
        }
    }

private:
    void sendChunk(const Chunk &chunk) {
        for (int i = 1;; i++) {  // 这里如果设置最大次数，达到最大次数时会导致对端卡死
            Client client(host, port);
            OperatorType op = CHUNK;
            client.send((const char *) &op, sizeof(op));
            client.send((const char *) &chunk.head, sizeof(chunk.head));
            client.send(chunk.data.get(), chunk.head.size);
            char res;
            client.recv(&res, 1);
            if (res == 1) {
                LOG("发送文件块 %d / %d 成功！\n", chunk.head.no, reader.getCount());
                break;
            } else {
                LOG("发送文件块 %d / %d 失败，第%d次尝试。\n", chunk.head.no, reader.getCount(), i);
            }
        }
    }

    void sendChunkWorker() {
        while (true) {
            Chunk chunk = reader.getNext();
            if (chunk.head.size == 0) {
                break;
            }
            sendChunk(chunk);
        }
    }

    void sendCollect() {
        Client client(host, port);
        OperatorType op = COLLECT;
        Collect collect{
                .count = reader.getCount(),
        };
        client.send((const char *) &op, sizeof(op));
        client.send((const char *) &collect, sizeof(collect));
    }

private:
    FileReader reader;
    const char *host;
    int port;
};
